/* 
* Copyright (C) 2010, John Asmuth

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
* 
*  $Revision$
*  $Date$
*  $Author$
*  $HeadURL$
* 
*/

package rlglue

type LocalExperiment struct {
	agent         Agent
	env           Environment
	lastAction    Action
	totalSteps    int32
	totalEpisodes int32
	totalReward   float64
}

func NewLocalExperiment(agent Agent, env Environment) (le *LocalExperiment) {
	le = &LocalExperiment{agent, env, nil, 0, 0, 0}
	return
}

func (ex *LocalExperiment) Init() (taskString string, err error) {
	taskString = ex.env.EnvInit()
	ex.agent.AgentInit(taskString)
	return
}

func (ex *LocalExperiment) Start() (obs Observation, act Action, err error) {
	ex.totalEpisodes++
	ex.totalSteps++
	obs = ex.env.EnvStart()
	act = ex.agent.AgentStart(obs)
	ex.lastAction = act
	return
}

func (ex *LocalExperiment) Step() (obs Observation, act Action, r float64, t bool, err error) {
	ex.totalSteps++
	obs, r, t = ex.env.EnvStep(ex.lastAction)
	if t {
		ex.agent.AgentEnd(r)
	} else {
		act = ex.agent.AgentStep(r, obs)
		ex.lastAction = act
	}
	ex.totalReward += r
	return
}

func (ex *LocalExperiment) Cleanup() (err error) {
	ex.agent.AgentCleanup()
	ex.env.EnvCleanup()
	return
}

func (ex *LocalExperiment) AgentMessage(message string) (reply string, err error) {
	reply = ex.agent.AgentMessage(message)
	return
}

func (ex *LocalExperiment) EnvironmentMessage(message string) (reply string, err error) {
	reply = ex.env.EnvMessage(message)
	return
}

func (ex *LocalExperiment) Return() (r float64, err error) {
	r = ex.totalReward
	return
}

func (ex *LocalExperiment) NumSteps() (numSteps int32, err error) {
	numSteps = ex.totalSteps
	return
}

func (ex *LocalExperiment) NumEpisodes() (numEpisodes int32, err error) {
	numEpisodes = ex.totalEpisodes
	return
}

func (ex *LocalExperiment) Episode(numSteps int32) (exitStatus int32, err error) {
	ex.Start()
	numSteps--
	for numSteps > 0 {
		_, _, _, t, _ := ex.Step()
		if t {
			break
		}
		numSteps--
	}
	return
}
